from flask import Flask, render_template, request, Response, stream_with_context, jsonify, url_for, send_from_directory
import requests
import json
import time
import os
import markdown
import uuid
from werkzeug.utils import secure_filename
from dotenv import load_dotenv
from flask_uploads import UploadSet, configure_uploads, IMAGES

# 加载环境变量
load_dotenv()

app = Flask(__name__)

# 配置上传文件存储位置
app.config['UPLOADED_PHOTOS_DEST'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制上传文件大小为16MB
app.config['SECRET_KEY'] = os.getenv("SECRET_KEY", "dev-key-for-chatbot")

# 确保上传目录存在
os.makedirs(app.config['UPLOADED_PHOTOS_DEST'], exist_ok=True)

# 设置允许上传的文件类型（扩展IMAGES类型以支持更多图片格式）
IMAGES = tuple('jpg jpe jpeg png gif svg webp'.split())
photos = UploadSet('photos', IMAGES)
configure_uploads(app, photos)

# 从环境变量获取API配置，如果不存在则使用默认值或明确的值
API_KEY = os.getenv("API_KEY", "sk-5IyouQJwhbbY2KRHJ2YtR9uJ2HEjGgH1JgG27f2PjDIA0ha4")
API_URL = os.getenv("API_URL", "https://apinexus.net/v1/chat/completions")
MODEL_ID = os.getenv("MODEL_ID", "gpt-4o-all")

# 存储对话历史
conversation_history = []

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    """提供上传文件的访问路径"""
    return send_from_directory(app.config['UPLOADED_PHOTOS_DEST'], filename)

@app.route('/upload', methods=['POST'])
def upload_file():
    """处理文件上传请求"""
    if 'file' not in request.files:
        return jsonify({'error': '没有文件'}), 400
    
    file = request.files['file']
    
    if file.filename == '':
        return jsonify({'error': '未选择文件'}), 400
    
    if file:
        try:
            # 生成安全的文件名并保存
            filename = secure_filename(file.filename)
            # 添加随机字符串防止文件名冲突
            name, ext = os.path.splitext(filename)
            unique_filename = f"{name}_{uuid.uuid4().hex[:8]}{ext}"
            
            filepath = photos.save(file, name=unique_filename)
            file_url = url_for('uploaded_file', filename=unique_filename)
            
            # 根据文件MIME类型返回不同的信息
            mime_type = file.mimetype
            if mime_type.startswith('image/'):
                file_type = 'image'
            else:
                file_type = 'file'
            
            return jsonify({
                'success': True,
                'file_url': file_url,
                'file_type': file_type,
                'filename': file.filename
            })
        except Exception as e:
            return jsonify({'error': f'上传失败: {str(e)}'}), 500
    
    return jsonify({'error': '上传失败'}), 400

@app.route('/chat', methods=['POST'])
def chat():
    # 获取用户消息和可能的图片URL
    user_message = request.form.get('message', '')
    image_url = request.form.get('image_url', '')
    
    if not user_message.strip() and not image_url:
        return Response("请输入有效的消息或上传图片", mimetype='text/plain')
    
    # 构建API请求
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    # 准备请求内容，支持多模态
    messages = []
    
    # 添加历史记录（如果需要）
    for msg in conversation_history:
        messages.append(msg)
    
    # 构建当前用户消息
    content = []
    
    # 如果有文本消息，添加到内容中
    if user_message.strip():
        content.append({
            "type": "text", 
            "text": user_message
        })
    
    # 如果有图片URL，添加到内容中
    if image_url:
        # 转换为绝对URL（如果是相对URL）
        if image_url.startswith('/'):
            image_url = request.host_url.rstrip('/') + image_url
        
        content.append({
            "type": "image_url",
            "image_url": {"url": image_url}
        })
    
    # 添加到消息列表
    user_message_obj = {"role": "user", "content": content if len(content) > 1 else content[0]["text"]}
    messages.append(user_message_obj)
    
    # 将当前用户消息添加到历史记录
    conversation_history.append(user_message_obj)
    
    data = {
        "model": MODEL_ID,
        "messages": messages,
        "stream": True  # 启用流式传输
    }
    
    # 定义流式传输的生成器函数
    def generate():
        try:
            response = requests.post(
                API_URL,
                headers=headers,
                data=json.dumps(data),
                stream=True,
                timeout=30  # 设置超时时间
            )
            
            if response.status_code != 200:
                error_msg = f"API返回错误 (状态码: {response.status_code})"
                try:
                    error_json = response.json()
                    if 'error' in error_json:
                        error_msg += f": {error_json['error'].get('message', '')}"
                except:
                    pass
                yield error_msg
                return
                
            # 收集完整的回复
            full_content = ""
            
            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data:'):
                        if line == 'data: [DONE]':
                            break
                        
                        try:
                            if line[5:].strip():  # 确保data:后面有内容
                                json_str = line[5:].strip()  # 去掉'data:'前缀和可能的空格
                                json_obj = json.loads(json_str)
                                delta = json_obj.get('choices', [{}])[0].get('delta', {})
                                content = delta.get('content', '')
                                if content:
                                    full_content += content
                                    yield content
                        except json.JSONDecodeError as e:
                            print(f"JSON解析错误: {e}, 行内容: {line}")
                            continue
                        except Exception as e:
                            print(f"处理流时出错: {e}")
                            yield f"处理响应时出错: {str(e)}"
                            return
            
            # 将AI回复添加到历史记录
            if full_content:
                conversation_history.append({"role": "assistant", "content": full_content})
                
                # 保持历史记录在合理范围内（例如最多保留10轮对话）
                if len(conversation_history) > 20:
                    conversation_history.pop(0)
                    conversation_history.pop(0)
            
        except requests.exceptions.Timeout:
            yield "请求超时，请稍后再试。"
        except requests.exceptions.ConnectionError:
            yield "连接服务器失败，请检查网络连接。"
        except Exception as e:
            print(f"请求出错: {e}")
            yield f"连接到API时出错: {str(e)}"
            return
    
    return Response(stream_with_context(generate()), mimetype='text/event-stream')

# 添加模拟聊天功能，以防API连接失败
@app.route('/chat/mock', methods=['POST'])
def mock_chat():
    user_message = request.form.get('message', '')
    image_url = request.form.get('image_url', '')
    
    def generate_mock_response():
        # 模拟回复，每个词之间有短暂延迟，模拟流式传输
        if image_url:
            response = f"我看到你上传了一张图片（{image_url}）并说：\"{user_message}\"。这是一个模拟回复，因为实际API可能暂时无法连接。我们正在一个个字符地流式传输这个回复，以展示流式传输的效果。\n\n这是一些**Markdown**格式的文本：\n\n# 标题1\n## 标题2\n\n- 列表项1\n- 列表项2\n\n```python\nprint('代码块示例')\n```"
        else:
            response = f"你好！我收到了你的消息: \"{user_message}\"。这是一个模拟回复，因为实际API可能暂时无法连接。我们正在一个个字符地流式传输这个回复，以展示流式传输的效果。\n\n这是一些**Markdown**格式的文本：\n\n# 标题1\n## 标题2\n\n- 列表项1\n- 列表项2\n\n```python\nprint('代码块示例')\n```"
        
        for word in response:
            yield word
            time.sleep(0.02)  # 模拟打字延迟
    
    return Response(stream_with_context(generate_mock_response()), mimetype='text/event-stream')

@app.route('/clear-history', methods=['POST'])
def clear_history():
    """清除对话历史"""
    global conversation_history
    conversation_history = []
    return jsonify({'success': True, 'message': '对话历史已清除'})

if __name__ == '__main__':
    app.run(debug=True) 